
import draw_curve
import torch
from torch import nn

#! 注意，对于卷积网络，输入图片的尺寸是固定的，动态调整卷积层池化层的步长自适应图像尺寸并不是一个好主意。
#! 相比之下，对于小图片，直接长宽等比放大，然后截取固定窗口识别；
#! 对于大图片，先高斯降采样进行识别，识别到物体，再从中截取原分辨率的图像进一步识别，这个思路更合理。 
#! 所以卷积层，最后如果有全连接层，要先计算好Flatten之后的尺寸。

"""
下面介绍一些torch里自带的层。每层的输入输出都是张量的形式。
零维张量：标量，
一维张量：矢量，
二维张量：矩阵……

nn.Module              # 就是神经网络的模块。（如：nn.Linear，nn.ReLU等）
nn.Linear(20, 256)     # 输入矢量，输出矢量。这里输入矢量长度为20，输出矢量长度为256。j本质就是一个矩阵线性变换(可能会加上偏置矢量)。
nn.ReLU()              # 输入什么尺寸的张量，就输出什么尺寸的张量。（对张亮的每个元素进行ReLU运算）
nn.Sequential( *args ) # args列表是nn.Module,负责将所有模块串联起来.前向传播的时候按顺序调用，上一层输出作为下一层输入。
                       # nn.Sequential本质也是nn.Module，相当于将整个网络都当做一个大的模块。

下面卷积层的模块比较复杂一些。  
nn.Conv2d  # 卷积核模块。                       
    参数说明
    in_channels：输入数据的通道数。
    out_channels：输出数据的通道数，每个输出通道对应一个不同的卷积核，用于提取不同的特征。
    kernel_size：卷积核的大小，可以是一个整数（表示正方形的卷积核），也可以是一个元组（如 (3, 4) 表示 3x4 的卷积核）。
    stride：卷积核在输入数据上滑动的步长。同样可以是一个整数或一个元组。用于下采样（缩小数据尺寸）。
    padding：填充操作的参数(零填充)。可以是一个整数，表示在输入数据的四周添加相同数量的行/列；也可以是一个元组，分别表示在上下和左右添加的行数/列数。
    设 输入nw*nh，核kw*kh，步长sw*sh，填充pw*ph，步长sw*sh，则输出 = (floor)[(nw-kw+pw+1)/sw] * (floor)[(nh-kh+ph+1)/sh] 。

    实现原理：
        设ci个输入通道，设co个输出通道，则有co个 ci通道的卷积核（每个通道的核都不一样，共co*ci个核）。
        先将输入通道与对应的卷积核通道单独做卷积运算，然后按元素相加，合并成一个通道，然后用co个这样的输出通道。 
        当卷积核是1*1的时候，主要用来改变通道数（压缩通道）
nn.AvgPool2d  # 最大池化层
nn.MaxPool2d  # 平均池化层
    参数说明：
    kernel_size：卷积核的大小，可以是一个整数（表示正方形的卷积核），也可以是一个元组（如 (3, 4) 表示 3x4 的卷积核）。
    stride：卷积核在输入数据上滑动的步长。同样可以是一个整数或一个元组。用于下采样（缩小数据尺寸）。
    功能：
        池化层主要用来改变输出大小（不改变输出通道数）
nn.Flatten    # 将张量转为矢量。
    在深度学习中，通常在卷积层或其他多维数据处理层之后使用，用于将提取到的特征图等多维数据结构整理成适合全连接层处理的格式。
    工作原理
        假设输入张量的形状为(batch_size, channels, height, width)（batch_size是批次大小，channels是通道数，height和width是特征图的高度和宽度），
        nn.Flatten会将除了batch_size之外的维度展平为一个维度。展平后的张量形状为(batch_size, channels * height * width)。
        训练好模型之后，实际应用的时候，batch_size通常为1.这时展平后的张量就是一维矢量。
    
"""



class Sequential(nn.Module):
    """ 这里复刻一下 nn.Sequential( *args ) ，理论上是可以替换使用的，但不建议。 能用库函数还是尽量用库函数。 """

    def __init__(self, *args):
        super().__init__()
        for idx, module in enumerate(args):
            # 这里，module是Module子类的一个实例。我们把它保存在'Module'类的成员
            # 变量_modules中。_module的类型是OrderedDict，有序字典
            self._modules[str(idx)] = module

    def forward(self, X):
        # OrderedDict保证了按照成员添加的顺序遍历它们
        for block in self._modules.values():
            if block is not None:
                X = block(X)
        return X



class Module_MLP_3(nn.Module):
    """_summary_  MLP : 多层感知机模块（Multilayer Perceptron）.
    这里定义3层多层感知机：一个输入层，一个隐藏层，一个输出层。
    隐藏层使用ReLU激活函数。
    """

    def __init__(self,in_size=20,out_size=10,hidden_size=256):
        """_summary_ MLP : 多层感知机模块（Multilayer Perceptron）.
    这里定义3层多层感知机：一个输入层，一个隐藏层，一个输出层。
    隐藏层使用ReLU激活函数。

        Args:
            in_size (int): _description_. 输入矢量长度
            out_size (int): _description_. 输出矢量长度
            hidden_size (int): _description_. 隐藏矢量长度，如果模型比较复杂，拟合效果不是很好，可以通过增加隐藏矢量长度增加模型复杂度。
        """

        # 调用MLP的父类Module的构造函数来执行必要的初始化。
        # 这样，在类实例化时也可以指定其他函数参数，例如模型参数params（稍后将介绍）
        super().__init__()
        # 输入矢量变换到隐藏矢量的矩阵
        self.in_hidden = nn.Linear(in_size,hidden_size)  
        # 隐藏矢量变换到输出矢量的矩阵
        self.hidden_out = nn.Linear(hidden_size, out_size)  

    # 定义模型的前向传播，即如何根据输入X返回所需的模型输出
    def forward(self, X):
        # 注意，这里我们使用ReLU的函数版本，其在nn.functional模块中定义。
        return self.out(nn.functional.relu(self.hidden(X)))



class Conv2D(nn.Module):
    """_summary_ 不要用自定义的卷积层，for i in range 实现效率太低。尽量用库。
        这里只是简单演示一下卷积层原理。
    Args:
        nn (_type_): _description_
    """

    def __init__(self, kernel_size):
        super().__init__()
        self.weight = nn.Parameter(torch.rand(kernel_size))
        self.bias = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return self.corr2d(x, self.weight) + self.bias

    def __corr2d__(self,X, K): 
        """_summary_ 计算二维互相关运算

        Args:
            X (矩阵): _description_ 原始图片
            K (矩阵): _description_ 卷积核
        Returns:
            矩阵: _description_  卷积计算的结果
        """
        h, w = K.shape  # 获取卷积核的高宽
        Y = torch.zeros((X.shape[0] - h + 1, X.shape[1] - w + 1))
        for i in range(Y.shape[0]):
            for j in range(Y.shape[1]):
                Y[i, j] = (X[i:i + h, j:j + w] * K).sum()
                # X[i:i + h, j:j + w] 矩阵切片。然后按元素与K相乘，得到新矩阵，然后对新矩阵所有元素求和。
        return Y


def save_net(net ,filename='output/net.params'):
    """_summary_

    Args:
        net (_type_): _description_ 保存网络模型。net是nn.modules 或者 nn.Sequential的变量
        filename (str, optional): 模型参数的保存路径
    """
    torch.save( net.state_dict(),filename)

def load_net(net ,filename='output/net.params'):
    """_summary_  使用方法： net = load_net(net ,filename) 。 

    Args:
        net (_type_): _description_ 保存网络模型。net是nn.modules 或者 nn.Sequential的变量
        filename (str, optional): 模型参数的保存路径

    Returns:
        _type_: _description_  返回加载好的网络模型。
    """
    net.load_state_dict(torch.load(filename,weights_only=True))
    #net.eval() # 这里将网络设为评估模式。
    return net





def print_net(net,is_show_para = False):
    """_summary_ 打印网络的结构。

    Args:
        net (_type_): _description_
        is_show_para (bool, optional): _description_. Defaults to False 是否显示权重参数，默认不显示.
    """
    for name, param in net.named_parameters(): # name是参数的名称，param是对应的参数张量，param.size()可以打印出参数张量的大小。
        print(name, param.size())
        if is_show_para:
            print(param)

def print_net_calculation_process(net,input_data,is_print_temp_shape = True,is_print_temp_data = False):
    """_summary_ 打印网络模型计算的中间过程。

    Args:
        net (_type_): _description_ 网络模型
        input_data (_type_): _description_ 输入张量。
        is_print_temp_shape (bool, optional): _description_. Defaults to True. 是否打印各层输出的形状。
        is_print_temp_data (bool, optional): _description_. Defaults to False. 是否打印各层输出的参数。
    Returns:
        _type_: _description_  返回神经网络计算的结果。
    """
    if is_print_temp_shape:
        print("\t input shape: \t",input_data.shape)
    if is_print_temp_data:
        print(input_data)

    for layer in net:
        input_data = layer(input_data)  # 上一层的输出结果作为下一层的输入。
        if is_print_temp_shape:
            print(layer.__class__.__name__,"\t output shape: \t",input_data.shape)
        if is_print_temp_data:
            print(input_data)
    return input_data

def test_save_net():
    X= torch.randn(2,20)
    nn.modules 
    net = nn.Sequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10))
    a = nn.Sequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10)) 
    # 第二次生成的网络参数会随机初始化，与第一个网络不一样。

    save_net(net)
    print(net)
    print(a)
    print(net(X))
    print(a(X))

    a = load_net(net)
    print(a)
    print(a(X))   
    print("判断加载前后a网络的输出即可判断是否保存的网络")



if __name__ == "__main__":


    #test_save_net()
 
    net = nn.Sequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10))
    #print_net(net)
    X= torch.randn(2,20)
    print_net_calculation_process(net,X,True)
















